package com.kotodama.infrastructure.tenant.connection;

import org.hibernate.engine.jdbc.connections.spi.AbstractDataSourceBasedMultiTenantConnectionProviderImpl;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 多租户连接提供者实现
 * 基于租户ID提供不同的数据源连接
 */
@Component
public class MultiTenantConnectionProviderImpl extends AbstractDataSourceBasedMultiTenantConnectionProviderImpl {

    private static final long serialVersionUID = 1L;
    private static final String DEFAULT_TENANT = "default";
    private final Map<String, DataSource> dataSourceMap = new ConcurrentHashMap<>();
    private DataSource defaultDataSource;

    public MultiTenantConnectionProviderImpl(DataSource defaultDataSource) {
        this.defaultDataSource = defaultDataSource;
        // 将默认数据源注册为默认租户
        dataSourceMap.put(DEFAULT_TENANT, defaultDataSource);
    }

    @Override
    protected DataSource selectAnyDataSource() {
        return defaultDataSource;
    }

    @Override
    protected DataSource selectDataSource(Object tenantIdentifier) {
        // 根据租户标识符选择数据源
        String tenantId = (String) tenantIdentifier;
        DataSource dataSource = dataSourceMap.get(tenantId);
        if (dataSource == null) {
            // 如果没有找到对应的数据源，使用默认数据源
            dataSource = defaultDataSource;
        }
        return dataSource;
    }

    /**
     * 注册租户数据源
     * 
     * @param tenantId   租户ID
     * @param dataSource 数据源
     */
    public void registerDataSource(String tenantId, DataSource dataSource) {
        dataSourceMap.put(tenantId, dataSource);
    }

    /**
     * 移除租户数据源
     * 
     * @param tenantId 租户ID
     */
    public void removeDataSource(String tenantId) {
        if (!DEFAULT_TENANT.equals(tenantId)) {
            dataSourceMap.remove(tenantId);
        }
    }

    /**
     * 获取所有已注册的租户ID
     * 
     * @return 租户ID集合
     */
    public java.util.Set<String> getRegisteredTenants() {
        return dataSourceMap.keySet();
    }
}